import json, time, re
from typing import List, Dict, Tuple, Any
from utils.logging_config import setup_logger
# from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
import torch

import httpx, openai
import requests

logger = setup_logger('Invoke LLM')

class CountryTranslate_Requests:
    def __init__(self):
        # 硅基流动
        self.url = "https://api.siliconflow.cn/v1/chat/completions"
        self.headers = {
            "Authorization": "Bearer sk-****",
            "Content-Type": "application/json"
        }

    def extract_evaluation(self, content) -> Dict[str, Any]:
        """
        从 ChatCompletionMessage 中提取并解析评估结果
        
        :param message: ChatCompletionMessage 对象
        :return: 解析后的评估字典
        """
        # pattern = r'```json(.*?)```'
        # match = re.search(pattern, content, re.DOTALL)
        # json_str = match.group(1)
        json_str = content

        if json_str:
            try:
                return json.loads(json_str)
            except json.JSONDecodeError as e:
                logger.info(f"JSON 解析错误: {e}")
                time.sleep(10)
                return {}
        else:
            logger.info("未找到有效的 JSON 内容")
            time.sleep(10)
            return {}
        
    def translate_country_lst(self, country_lst: List, country_code_dict: List[Dict]) -> List[Dict]:
        """
        使用LLM评估多轮对话质量
        """
        start=time.time()
        # 构建评估提示
        prompt = f"""
        您是一个专业的多语种语言翻译专家。以下列表中有{len(country_lst)}个元素，包含多语种国家名称、两位或三位国别编码。请对其进行翻译和转换，并与以下数据字典
        中{len(country_code_dict)}个国家的中文国名、英文国名进行对应。如果对应不上，请直接返回空值。
    
        ### 需要翻译的数据列表：
        {country_lst}
    
        ### 需要映射的内容
        {json.dumps(country_code_dict, indent=2, ensure_ascii=False)}    
        
        ### 输出格式要求（JSON）：
        [{{
          "原始国名": <str>,
          "中文国名": <str>,
          "英文国名": <str>
        }},
        ...]
        """
        
        # 调用LLM API
        try:
            payload = {
                # "model": "qwen3",
                # "model": "Qwen/Qwen3-14B",
                # "model": "Qwen/Qwen3-32B",
                "model": "Qwen/Qwen3-235B-A22B-Instruct-2507",
                "messages": [
                    {"role": "system", "content": "你是一个专业的多语种语言翻译专家"},
                    {"role": "user", "content": prompt}
                ],
                "extra_body": {"chat_template_kwargs": {"enable_thinking": False}},
                "temperature": 0.3,
                "stream": False
            }
            result = {}
            while result=={}:
                response = requests.post(self.url, json=payload, headers=self.headers)
                if response.status_code==200:
                    rslt = response.json()
                    # print(rslt["choices"][0]["message"]["content"])
                    # result = self.extract_evaluation(rslt["choices"][0]["message"]["content"])
                    result = rslt["choices"][0]["message"]["content"]
                    logger.info(result)
                    # del response
                    torch.cuda.empty_cache()
            end=time.time()
            logger.info(f'Invoke LLM takes {round(end-start)} seconds')
            return result
        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                torch.cuda.empty_cache()
                time.sleep(1)
                return self.process_batch(prompts)
            raise {"error": "Failed to parse LLM response"}

def extract_evaluation(content) -> Dict[str, Any]:
    """
    从 ChatCompletionMessage 中提取并解析评估结果
    
    :param message: ChatCompletionMessage 对象
    :return: 解析后的评估字典
    """
    # pattern = r'```json(.*?)```'
    # match = re.search(pattern, content, re.DOTALL)
    # json_str = match.group(1)
    json_str = content
    if json_str:
        try:
            return json.loads(json_str)
        except json.JSONDecodeError as e:
            logger.info(f"JSON 解析错误: {e}")
            time.sleep(10)
            return {}
    else:
        logger.info("未找到有效的 JSON 内容")
        time.sleep(10)
        return {}